from engine.validator import BaseValidator
from src.evaluate import BLEU
import torch
from tqdm import tqdm


class TranslationValidator(BaseValidator):
    """机器翻译模型验证类，继承自BaseValidator"""
    def __init__(self, model, dataloader, criterion, device, tokenizer):
        """
        param model: 机器翻译Transformer模型
        param dataloader: 验证数据集加载器
        param criterion: 损失函数
        param device: 运行设备
        Param tokenizer: 词汇表
        """
        super().__init__(model, dataloader, criterion, device)
        self.bleu = BLEU()
        self.tokenizer = tokenizer  # 用于解码预测结果

    def generate_masks(self, src, tgt):
        """生成源语言和目标语言掩码"""
        src_mask = (src != 0).unsqueeze(1).unsqueeze(2)
        seq_len = tgt.size(1)
        nopeak_mask = torch.tril(torch.ones(seq_len, seq_len, device=tgt.device)).bool()
        tgt_mask = ((tgt != 0).unsqueeze(1).unsqueeze(2)) & nopeak_mask.unsqueeze(0)
        return src_mask, tgt_mask



    def generate_translation(self, src):
        """
        通过Transformer生成翻译文本
        :param src: 输入源句子张量
        :return: 生成的翻译文本张量
        """
        self.model.eval()
        total_loss = 0
        reference, translations = [], []
        with torch.no_grad():
            translated_tokens = self.model.generate(src)
        return translated_tokens

    def validate(self):
        self.model.eval()
        total_loss = 0
        references, translations = [], []

        with torch.no_grad():
            pbar = tqdm(self.dataloader, total=len(self.dataloader), desc=f"Validating", leave=True)
            for batch in pbar:
                src, tgt = batch['src'].to(self.device), batch['tgt'].to(self.device)
                src_mask, tgt_mask = self.generate_masks(src, tgt)

                # 预测翻译（传递 self.tgt_vocab）
                output = self.model.generate(src, self.tokenizer)

                # 计算损失
                seq_len = tgt[:, :-1].size(1)
                logits = self.model(src, tgt[:, :-1], src_mask, tgt_mask[:, :, :seq_len, :seq_len])
                loss = self.criterion(logits.view(-1, logits.size(-1)), tgt[:, 1:].reshape(-1))
                total_loss += loss.item()
                pbar.set_postfix({'val_loss': loss.item()})

                # 记录 BLEU 计算所需的参考文本和翻译文本
                references.extend(tgt.cpu().numpy().tolist())
                translations.extend(output.cpu().numpy().tolist())

        avg_loss = total_loss / len(self.dataloader)

        # **计算 BLEU 分数**
        bleu_calculator = BLEU()
        bleu_score = sum(bleu_calculator.compute(ref, trans) for ref, trans in zip(references, translations)) / len(
            references)

        return avg_loss, bleu_score

